#ifndef NPU_TRITON_ADD_KERNEL_H
#define NPU_TRITON_ADD_KERNEL_H

#include <torch/torch.h>
#include <limits>
#include <acl/acl.h>
#include "kernel_launchers.h"
#include <torch_npu/csrc/aten/NPUNativeFunctions.h>
#include <torch_npu/csrc/core/npu/NPUStream.h>

namespace TritonTorch {

/**
 * @brief Launch the add kernel on NPU device
 * 
 * @param x Input tensor x (must be on NPU, float32, contiguous)
 * @param y Input tensor y (must be on NPU, float32, contiguous)
 * @param nElements Number of elements to process
 * @param blockSize Block size for kernel launch
 * @return torch::Tensor Output tensor (same shape as input)
 */
torch::Tensor launch_add_kernel(const torch::Tensor& x,
                                const torch::Tensor& y,
                                int64_t nElements,
                                int32_t blockSize, 
                                int32_t gridX, int32_t gridY = 1, int32_t gridZ = 1);

} // namespace TritonTorch

#endif // NPU_TRITON_ADD_KERNEL_H

